import torch
import numpy as np
import pandas as pd

from torch import nn
import torch.nn.functional as F
from torch import optim
import matplotlib.pyplot as plt
from torch.autograd import Variable

df_train_data = pd.read_csv('abalone_train.data')
df_test_data = pd.read_csv('abalone_test.data')

train_data_x = df_train_data.iloc[:, 1:10]
train_data_x = np.array(train_data_x, dtype=float)
train_data_x = torch.unsqueeze(torch.FloatTensor(train_data_x), dim=1)
print(train_data_x)

test_data_x = df_train_data.iloc[:, 1:10]
test_data_x = np.array(test_data_x, dtype=float)
test_data_x = torch.unsqueeze(torch.FloatTensor(test_data_x), dim=1)
print(test_data_x)

df_train_data['Sex Code'] = df_train_data['Sex'].map({'M': 0, 'F': 1, 'I': 2})
df_test_data['Sex Code'] = df_test_data['Sex'].map({'M': 0, 'F': 1, 'I': 2})

train_data_y = df_train_data['Sex Code']
train_data_y = np.array(train_data_y, dtype=int)
train_data_y = torch.unsqueeze(torch.FloatTensor(train_data_y), dim=1)
print(train_data_y)


test_data_y = df_train_data['Sex Code']
test_data_y = np.array(test_data_y, dtype=int)
test_data_y = torch.unsqueeze(torch.FloatTensor(test_data_y), dim=1)
print(test_data_y)

# Step 2:============================定义模型===================
# 定义一个类，继承自 torch.nn.Module，torch.nn.Module是callable的类
# 在整个类里面重新定义一个标准的BP全连接神经网络，网络一共是四层，
# 层数定义：784, 400， 200， 100， 10
# 其中输入层784个节点，输出层是10个节点，分别代表10个数字，其他的层都是隐藏层。
# 我们使用了Relu的激活函数，而不是sigmoid激活函数
# 整个子类需要重写forward函数，
class BPNNModel(torch.nn.Module):
    def __init__(self):
        # 调用父类的初始化函数，必须要的
        super(BPNNModel, self).__init__()

        # 创建四个Sequential对象，Sequential是一个时序容器，将里面的小的模型按照序列建立网络
        self.layer1 = nn.Sequential(nn.Linear(784, 300), nn.ReLU())
        self.layer2 = nn.Sequential(nn.Linear(300, 150), nn.ReLU())
        self.layer3 = nn.Sequential(nn.Linear(150, 45), nn.ReLU())
        self.layer4 = nn.Sequential(nn.Linear(45, 3))

    def forward(self, x):
        # 每一个时序容器都是callable的，因此用法也是一样。
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        return x


# 创建和实例化一个整个模型类的对象
model = BPNNModel()
# 打印出整个模型
print(model)

# Step 3:============================定义损失函数和优化器===================
# 定义 loss 函数，这里用的是交叉熵损失函数(Cross Entropy)，这种损失函数之前博文也讲过的。
criterion = nn.CrossEntropyLoss()
# 我们优先使用随机梯度下降，lr是学习率: 0.1
optimizer = torch.optim.SGD(model.parameters(), 1e-1)

# Step 4:============================开始训练网络===================
# 为了实时观测效果，我们每一次迭代完数据后都会，用模型在测试数据上跑一次，看看此时迭代中模型的效果。
# 用数组保存每一轮迭代中，训练的损失值和精确度，也是为了通过画图展示出来。
train_losses = []
train_acces = []
# 用数组保存每一轮迭代中，在测试数据上测试的损失值和精确度，也是为了通过画图展示出来。
eval_losses = []
eval_acces = []

for e in range(20):

    # 4.1==========================训练模式==========================
    train_loss = 0
    train_acc = 0
    model.train()  # 将模型改为训练模式

    # 每次迭代都是处理一个小批量的数据，batch_size是64
    for im, label in train_data_x:
        im = Variable(im)
        label = Variable(label)

        # 计算前向传播，并且得到损失函数的值
        out = model(im)
        loss = criterion(out, label)

        # 反向传播，记得要把上一次的梯度清0，反向传播，并且step更新相应的参数。
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # 记录误差
        train_loss += loss.item()

        # 计算分类的准确率
        _, pred = out.max(1)
        num_correct = (pred == label).sum().item()
        acc = num_correct / im.shape[0]
        train_acc += acc

    train_losses.append(train_loss / len(train_data_x))
    train_acces.append(train_acc / len(train_data_x))

    # 4.2==========================每次进行完一个训练迭代，就去测试一把看看此时的效果==========================
    # 在测试集上检验效果
    eval_loss = 0
    eval_acc = 0
    model.eval()  # 将模型改为预测模式

    # 每次迭代都是处理一个小批量的数据，batch_size是128
    for im, label in test_data_x:
        im = Variable(im)  # torch中训练需要将其封装即Variable，此处封装像素即784
        label = Variable(label)  # 此处为标签

        out = model(im)  # 经网络输出的结果
        loss = criterion(out, label)  # 得到误差

        # 记录误差
        eval_loss += loss.item()

        # 记录准确率
        _, pred = out.max(1)  # 得到出现最大值的位置，也就是预测得到的数即0—9
        num_correct = (pred == label).sum().item()  # 判断是否预测正确
        acc = num_correct / im.shape[0]  # 计算准确率
        eval_acc += acc

    eval_losses.append(eval_loss / len(test_data_x))
    eval_acces.append(eval_acc / len(test_data_x))
    print('epoch: {}, Train Loss: {:.6f}, Train Acc: {:.6f}, Eval Loss: {:.6f}, Eval Acc: {:.6f}'
          .format(e, train_loss / len(test_data_x), train_acc / len(test_data_x),
                  eval_loss / len(test_data_x), eval_acc / len(test_data_x)))

plt.title('train loss')
plt.plot(np.arange(len(train_losses)), train_losses)
plt.plot(np.arange(len(train_acces)), train_acces)
plt.title('train acc')
plt.plot(np.arange(len(eval_losses)), eval_losses)
plt.title('test loss')
plt.plot(np.arange(len(eval_acces)), eval_acces)
plt.title('test acc')
plt.show()
